!##############################################################################
! MODULE globals
!
! This code defines GLOBALS AND FUNCTIONS to solve life-cycle problem 
! with couples + singles, CRRA, det. medical exp. during retirement
!
! Author: Annika Bacher (annika.bacher@bi.no)
!          
!
! This version: May 2024
!
!##############################################################################


module globals_ret

        use toolbox_own
        use globals

        implicit none
        
        ! 1) length of retirement

        integer, parameter :: JR = 20, retage = 65
        
        !  ********  gender specific ********
         
        ! 2) female medical expenditures

        real*8 :: m_thetaf(naa), m_efff(JR),m_yf, m_eff1f, m_eff2f, m_eff3f
        
        ! 3) male medical expenditures
        
        real*8 :: m_thetam(naa), m_effm(JR), m_ym, m_eff1m, m_eff2m, m_eff3m

        ! 4) pension payments, equivalence scales, survival prob.

        real*8, dimension(nzz,naa) :: pen_c, pen_f, pen_m
        
        real*8, dimension(JR) :: nu_retf , nu_retm , nu_retc, psim, psif
    
        
        ! 5) Value Functions and Optimization

        real*8, dimension(nrisk,ncc,nzz,naa,JR)         :: Voptret_f, Voptret_m, cash_optret_m, cash_optret_f
                                                                                      
        real*8, dimension(ncc,nzz,naa,JR)               :: VFret_f, VFret_m

        real*8, dimension(nrisk,ncc,nzz,naa,naa,JR)     :: Voptret_c, cash_optret_c, Voptret_ci, Voptret_cp        

        real*8, dimension(ncc,nzz,naa,naa,JR)           :: VFret_c, VFret_ci, VFret_cp

        
        ! 6) Policy Functions when single
        
        real*8, dimension(nrisk,ncc,nzz,naa,JR) :: coptret_f, coptret_m
        
        real*8, dimension(ncc,nzz,naa,JR)       :: sharepolret_f, sharepolret_m,riskypolret_f, riskypolret_m, &
                                               savepolret_f, savepolret_m,cpolret_f, cpolret_m, &
                                               cprimepolret_f, cprimepolret_m
                                             
          
        ! 7) Policy Functions when being in a couple                                         
        
         real*8, dimension(JR,nzz,naa,naa,ncc)      :: cpolret_c, savepolret_c, sharepolret_c, &
                                                      riskypolret_c, cprimepolret_c
                                                     

        real*8, dimension(JR,nzz,naa,naa,ncc,nrisk) :: coptret_c 
                                                           

        ! 11) Displaying the results
        
        real*8 :: age_ret(JR)
        
        ! 12) Continuation Value for working period
        
        ! continuation values by marital status 
        ! this is the object need as continuation value for working age
        
        real*8, dimension(ncc,nzz,naa)  :: cont_m , cont_f
        real*8, dimension(ncc,nzz,naa,naa) :: cont_c, cont_ci, cont_cp
        
        ! 13) human capital returns
        
        real*8, dimension(JR,nzz*ncc,nzz*ncc) :: distret_f_low, distret_f_high, distret_m_low, distret_m_high

   
!##############################################################################  

!                           INCLUDED FUNCTIONS  
  
!##############################################################################  
  
        
contains


    ! 1) function for optimal intertemporal decision, singles

        function valfun_sing_ret(cprime,share,delta,gam,y,theta,eff,beta,pen,c0, &
                                    vfs,nu,psi,SMF) result(valfres_ret)


        ! declare variables
    
            implicit none
    
            integer ::  i
    
            integer, dimension(nrr,2) ::  clowst, chighst 
            
            real*8, dimension(nrr,2) :: cashtmrst
            
            real*8  ::  gg_s, valfres_ret(2), SMF
    
            real*8  :: share, c0,cprime, consum, pen, nu
            
            real*8  :: theta, eff, gam, beta, inf, delta
    
            real*8  :: valf, y, exp_s,psi
    
            real*8, dimension(ncc) :: vfs
    

        ! define infinity
    
            inf = 900000000000000000000000000d0

        
        ! HH chooses cprime, specify cash-on-hand next period
        ! depends on aggregate state and return realization of risky asset
    
            do i=1,nrr
    
            cashtmrst(i,1) = pen+ ((1d0-share)*(1d0+r)+share*ard_boom(i))*cprime
            cashtmrst(i,2) = pen+ ((1d0-share)*(1d0+r)+share*ard_rec(i))*cprime
    
            clowst(i,1)    = max(cntarray(cgrid,ncc,cashtmrst(i,1)),1)
            chighst(i,1)   = min(clowst(i,1) + 1,ncc)
            
            clowst(i,2)    = max(cntarray(cgrid,ncc,cashtmrst(i,2)),1)
            chighst(i,2)   = min(clowst(i,2) + 1,ncc)

            
            if (clowst(i,1) == ncc) clowst(i,1)  = clowst(i,1)  -1 
            if (clowst(i,2) == ncc) clowst(i,2)  = clowst(i,2)  -1 

            end do

  
        ! specify the continuation value for single retirees, take expected value over realizations
        ! linear interpolation of value function
            gg_s = 0d0
                
            do i = 1,nrr
            gg_s = gg_s + ( (1d0-p_rec)*(pir_boom(i)*(vfs(clowst(i,1)) + &
                (cashtmrst(i,1)-cgrid(clowst(i,1)))*(vfs(chighst(i,1)) - &
                vfs(clowst(i,1)))/(cgrid(chighst(i,1)) - cgrid(clowst(i,1))))) + &
                p_rec*(pir_rec(i)*(vfs(clowst(i,2)) + &
                (cashtmrst(i,2)-cgrid(clowst(i,2)))*(vfs(chighst(i,2)) - &
                vfs(clowst(i,2)))/(cgrid(chighst(i,2)) - cgrid(clowst(i,2))))) )
            end do
        
        ! compute consumption
            if(share==0d0) then
            consum = c0-cprime-(eff*theta*y)
            else
            consum = c0-SMF-cprime-(eff*theta*y)
            end if
            
        ! impossible allocations ('consumption floor')

            if(consum<=0.0002d0) consum = 0.0002d0
          
        !  expected value tomorrow (pension income is deterministic)
                              
            exp_s=gg_s
          
        ! compute the value function
        
			! ROBUSTNESS: change 0.128d0 to 12.8d0
    
            valf = nu*((consum/nu)**(1d0-gam))/(1d0-gam) + &
                beta*psi*exp_s + beta*(1d0-psi)*(0.128d0*(delta+cprime)**(1d0-gam))/(1d0-gam)
            
            valfres_ret(1) = valf
            valfres_ret(2) = consum

            end function valfun_sing_ret
            

    ! 2) function for optimal intertemporal decision, couples

        function valfun_coup_ret(cprime,share,deltas,deltac,gami,gamp,yi,yp,thetai,thetap,effi,effp, &
                                        betac,betai, betap,pen,c0, &
                                vfc,vfi,vfp,vci,vcp,nu,psii,psip,SMF,wid_f,wid_m) result(val_results_ret)
        
        ! declare variables
        
            implicit none
            
            integer ::   ir
    
            integer, dimension(nrr,2) ::  clowst, chighst, clowst_f, chighst_f, clowst_m, chighst_m
            
            real*8, dimension(nrr,2)  ::  cashtmrst, cashtmrst_f, cashtmrst_m 
            
            real*8  :: gg_si, gg_sp,  gg_c, gg_vci, gg_vcp
    
            real*8  ::  share, c0,cprime, cprime_f, cprime_m, consum, pen,inf, val_results_ret(4) 
            
            real*8  ::  yi, yp, thetai, thetap, effi, effp, gami, gamp, betai, betap, betac, nu, deltas, deltac
    
            real*8  ::  valf, val_ci, val_cp, wid_f, wid_m
    
            real*8  ::  exp_c, psii, psip,exp_si, exp_sp, SMF, exp_vci, exp_vcp
            
            real*8, dimension(ncc) :: vfc, vfi, vfp, vci, vcp

         
        ! define infinity
    
            inf = 900000000000000000000000000d0


        ! HH chooses cprime    
        ! in case of widowhood, surviving spouse keeps fraction of assets
        ! depends on whether husband or wife dies
        
            cprime_f = cprime*0.6d0   
            cprime_m = cprime*0.7d0
            
       ! ROBUSTNESS: NO ASSET DROP AFTER DEATH OF ONE SPOUSE
           
            !cprime_f = cprime  
            !cprime_m = cprime
            
        ! specify cash-on-hand next period
        ! depends on aggregate state and return realization of risky asset
    
            do ir=1,nrr
            
            ! couples
            cashtmrst(ir,1) = pen + ((1d0-share)*(1d0+r)+share*ard_boom(ir))*cprime
            cashtmrst(ir,2) = pen + ((1d0-share)*(1d0+r)+share*ard_rec(ir))*cprime
            
            clowst(ir,1)    = max(cntarray(cgrid,ncc,cashtmrst(ir,1)),1)
            chighst(ir,1)   = min(clowst(ir,1) + 1,ncc)
            
            clowst(ir,2)    = max(cntarray(cgrid,ncc,cashtmrst(ir,2)),1)
            chighst(ir,2)   = min(clowst(ir,2) + 1,ncc)
            
            if (clowst(ir,1) == ncc) clowst(ir,1)  = clowst(ir,1)  -1 
            if (clowst(ir,2) == ncc) clowst(ir,2)  = clowst(ir,2)  -1
    
			! women in case of widowhood 
            cashtmrst_f(ir,1) = wid_f + ((1d0-share)*(1d0+r)+share*ard_boom(ir))*cprime_f
            cashtmrst_f(ir,2) = wid_f + ((1d0-share)*(1d0+r)+share*ard_rec(ir))*cprime_f
    
            clowst_f(ir,1)    = max(cntarray(cgrid,ncc,cashtmrst_f(ir,1)),1)
            chighst_f(ir,1)   = min(clowst_f(ir,1) + 1,ncc)
            
            clowst_f(ir,2)    = max(cntarray(cgrid,ncc,cashtmrst_f(ir,2)),1)
            chighst_f(ir,2)   = min(clowst_f(ir,2) + 1,ncc)
            
            if (clowst_f(ir,1) == ncc) clowst_f(ir,1)  = clowst_f(ir,1)  -1 
            if (clowst_f(ir,2) == ncc) clowst_f(ir,2)  = clowst_f(ir,2)  -1
            
			! men in case of widowhood
            cashtmrst_m(ir,1) = wid_m +((1d0-share)*(1d0+r)+share*ard_boom(ir))*cprime_m
            cashtmrst_m(ir,2) = wid_m +((1d0-share)*(1d0+r)+share*ard_rec(ir))*cprime_m
    
            clowst_m(ir,1)    = max(cntarray(cgrid,ncc,cashtmrst_m(ir,1)),1)
            chighst_m(ir,1)   = min(clowst_m(ir,1) + 1,ncc)
            
            clowst_m(ir,2)    = max(cntarray(cgrid,ncc,cashtmrst_m(ir,2)),1)
            chighst_m(ir,2)   = min(clowst_m(ir,2) + 1,ncc)
            
            if (clowst_m(ir,1) == ncc) clowst_m(ir,1)  = clowst_m(ir,1)  -1 
            if (clowst_m(ir,2) == ncc) clowst_m(ir,2)  = clowst_m(ir,2)  -1
            
            end do
            

        ! specify the continuation values, linear interpolation of value function
        
        ! Value of single, in case of widowhood (seperately for i and p)
        
            gg_si = 0d0
            gg_sp = 0d0
            gg_c  = 0d0
        
 
            do ir = 1,nrr
            gg_si = gg_si+ ((1d0-p_rec)*( pir_boom(ir)*(vfi(clowst_f(ir,1)) + &
               (cashtmrst_f(ir,1)-cgrid(clowst_f(ir,1)))*(vfi(chighst_f(ir,1)) - &
               vfi(clowst_f(ir,1)))/(cgrid(chighst_f(ir,1)) - cgrid(clowst_f(ir,1)))))+&
               p_rec*( pir_rec(ir)*(vfi(clowst_f(ir,2)) + &
               (cashtmrst_f(ir,2)-cgrid(clowst_f(ir,2)))*(vfi(chighst_f(ir,2)) - &
               vfi(clowst_f(ir,2)))/(cgrid(chighst_f(ir,2)) - cgrid(clowst_f(ir,2))))) )
               
               
            end do 
         
            do ir = 1,nrr
            gg_sp = gg_sp+ ((1d0-p_rec)*( pir_boom(ir)*(vfp(clowst_m(ir,1)) + &
               (cashtmrst_m(ir,1)-cgrid(clowst_m(ir,1)))*(vfp(chighst_m(ir,1)) - &
               vfp(clowst_m(ir,1)))/(cgrid(chighst_m(ir,1)) - cgrid(clowst_m(ir,1))))) + &
                 p_rec*( pir_rec(ir)*(vfp(clowst_m(ir,2)) + &
               (cashtmrst_m(ir,2)-cgrid(clowst_m(ir,2)))*(vfp(chighst_m(ir,2)) - &
               vfp(clowst_m(ir,2)))/(cgrid(chighst_m(ir,2)) - cgrid(clowst_m(ir,2))))) )
               
            end do 
        
     
        ! Value of couple
   
            do ir = 1,nrr
            gg_c = gg_c+ ((1d0-p_rec)*(pir_boom(ir)*(vfc(clowst(ir,1)) + &
                (cashtmrst(ir,1)-cgrid(clowst(ir,1)))*(vfc(chighst(ir,1)) - &
                vfc(clowst(ir,1)))/(cgrid(chighst(ir,1)) - cgrid(clowst(ir,1))))) + & 
                p_rec*(pir_rec(ir)*(vfc(clowst(ir,2)) + &
                (cashtmrst(ir,2)-cgrid(clowst(ir,2)))*(vfc(chighst(ir,2)) - &
                vfc(clowst(ir,2)))/(cgrid(chighst(ir,2)) - cgrid(clowst(ir,2))))) )  
            end do
            
            
        ! Value of individual in couple (separetly for i and p)

			gg_vci = 0d0
			gg_vcp = 0d0

            do ir = 1,nrr
            gg_vci = gg_vci+ ((1d0-p_rec)*(pir_boom(ir)*(vci(clowst(ir,1)) + &
               (cashtmrst(ir,1)-cgrid(clowst(ir,1)))*(vci(chighst(ir,1)) - &
               vci(clowst(ir,1)))/(cgrid(chighst(ir,1)) - cgrid(clowst(ir,1))))) + &
               p_rec*(pir_rec(ir)*(vci(clowst(ir,2)) + &
               (cashtmrst(ir,2)-cgrid(clowst(ir,2)))*(vci(chighst(ir,2)) - &
               vci(clowst(ir,2)))/(cgrid(chighst(ir,2)) - cgrid(clowst(ir,2))))) )
            end do 
            
            do ir = 1,nrr
            gg_vcp = gg_vcp+ ((1d0-p_rec)*(pir_boom(ir)*(vcp(clowst(ir,1)) + &
               (cashtmrst(ir,1)-cgrid(clowst(ir,1)))*(vcp(chighst(ir,1)) - &
               vcp(clowst(ir,1)))/(cgrid(chighst(ir,1)) - cgrid(clowst(ir,1))))) + &
               p_rec*(pir_rec(ir)*(vcp(clowst(ir,2)) + &
               (cashtmrst(ir,2)-cgrid(clowst(ir,2)))*(vcp(chighst(ir,2)) - &
               vcp(clowst(ir,2)))/(cgrid(chighst(ir,2)) - cgrid(clowst(ir,2))))) )
            end do 
            
        ! compute consumption
    
            if(share==0d0) then
            consum = c0-cprime-(1d0*yi*thetai*effi+1d0*yp*thetap*effp)     ! ROBUSTNESS: change 1d0 to 0.8d0
            else
            consum = c0-cprime-SMF-(1d0*yi*thetai*effi+1d0*yp*thetap*effp) ! ROBUSTNESS: change 1d0 to 0.8d0
            end if

        ! impossible allocations ('consumption floor')
        
            if(consum<=0.0002d0) consum = 0.0002d0
                    
        ! expected value tomorrow (pension income is deterministic)
        
            ! widowhood of i 
            exp_si=gg_si
            
            ! widowhood of p 
            exp_sp=gg_sp

            ! both spouses survive 
            exp_c=gg_c
            
            ! value of i in couple
            exp_vci = gg_vcp
            
            ! value of p in couple
            exp_vcp = gg_vcp
    
        
        ! ROBUSTNESS: change 0.128d0 to 12.8d0
            
        ! compute the value function and of individual being in THAT couple
            valf    = nu*((consum/nu)**(1d0-gami))/(1d0-gami) + & 
                        betac*psii*psip*exp_c +betac*psii*(1d0-psip)*exp_si+betac*psip*(1d0-psii)*exp_sp+ &  
                        betac*(1d0-psii)*(1d0-psip)*(0.128d0*(deltac+cprime)**(1d0-gami))/(1d0-gami)  
                    
                                      
            val_ci =  nu*((consum/nu)**(1d0-gami))/(1d0-gami) + betac*psii*psip*exp_vci + &
                    betac*psii*(1d0-psip)*exp_si + betac*(1d0-psii)*(1d0-psip)*(0.128d0*(deltas+cprime)**(1d0-gami))/(1d0-gami)
                                                
            
            val_cp =  nu*((consum/nu)**(1d0-gamp))/(1d0-gamp) + betac*psip*exp_vcp + &
                    betac*psip*(1d0-psii)*exp_sp + betac*(1d0-psii)*(1d0-psip)*(0.128d0*(deltas+cprime)**(1d0-gamp))/(1d0-gamp)
                             
        
            val_results_ret(1) = valf
            val_results_ret(2) = consum
            val_results_ret(3) = val_ci
            val_results_ret(4) = val_cp
            

        end function valfun_coup_ret
        
    ! 3) golden function (single, CD) - female
    
        function golden_sing_ret(aa,bb,share,delta,gam,y,theta,eff,beta, &
                                pen,c0,vfs,nu,psi,SMF) result(output_ret)
                                

            ! declare variables
            
            implicit none
    
            real*8  :: alpha1, alpha2, d, x1, x2, tol, output_ret(3), f1, f2
            
            real*8  :: aa, bb, delta
            
            integer ::  l
    
            real*8  :: share, c0, pen, valfres_ret(2),SMF
            
            real*8  :: y, theta, eff, gam, beta, nu, psi

            real*8, dimension(ncc) :: vfs

            
            ! find optimal value with golden section search
    
            tol     =   0.0000000000000001d0
            alpha1  =   (3d0-sqrt(5d0))/2d0
            alpha2  =   (sqrt(5d0)-1d0)/2d0
                d   =   bb-aa
                x1  =   aa+alpha1*d   
                x2  =   aa+alpha2*d   
                
                valfres_ret  =   &
                valfun_sing_ret(x1,share,delta,gam,y,theta,eff,beta,pen,c0,vfs,nu,psi,SMF)
                f1  = valfres_ret(1)
                valfres_ret  =   &
                valfun_sing_ret(x2,share,delta,gam,y,theta,eff,beta,pen,c0,vfs,nu,psi,SMF)
                f2  = valfres_ret(1)
                d   =   alpha1*alpha2*d
                
                if(f1==-900000000000000000000000000d0 .AND. share == 0d0) write(*,*) 'vf is infinity -- single', f2, l 
              
            do while(d>tol)
                       
                d = d*alpha2

            if(f2<f1)then

                x2 = x1
                x1 = x1-d
                f2 = f1
                valfres_ret  =  &
                valfun_sing_ret(x1,share,delta,gam,y,theta,eff,beta,pen,c0,vfs,nu,psi,SMF)
                f1 = valfres_ret(1)
 
            else

                x1 = x2
                x2 = x2+d
                f1 = f2
                valfres_ret  = &
                 valfun_sing_ret(x2,share,delta,gam,y,theta,eff,beta,pen,c0,vfs,nu,psi,SMF)
                f2 = valfres_ret(1)
   
            end if
            
            end do    

            if(f2>f1) then
        
            x1=x2
            f1=f2
        
            end if
            
            ! store results in array
 
            output_ret(1) = x1
            output_ret(2) = f1
            output_ret(3) = valfres_ret(2)

    end function golden_sing_ret

    ! 4) golden function (couple, CD)

        function golden_coup_ret(aa,bb,share,deltas,deltac,gami,gamp,yi,yp,thetai,thetap,effi,effp, &
                                betac, betai,betap,pen,c0, &
                                vfc,vfi,vfp,vci,vcp,nu,psii,psip,SMF,wid_f,wid_m) result(output_coup_ret)
                                
                                
            ! declare variables
            
            implicit none
    
            real*8  :: alpha1, alpha2, d, x1, x2, tol, output_coup_ret(5), f1, f2
            
            real*8  :: aa, bb, r_intertmp_ret(4),SMF
            
            integer :: l
    
            real*8  :: share, c0, pen, deltas, deltac
            
            real*8  :: yi, yp, thetai, thetap, effi, effp, gami, gamp,betai,betap,betac,nu
    
            real*8  :: psii, psip,wid_f,wid_m
            
            real*8, dimension(ncc) :: vfi, vfp, vci, vcp, vfc
    
                                        
            ! find optimal value with golden section search
    
             tol     =   0.0000000000000001d0
            alpha1  =   (3d0-sqrt(5d0))/2d0
            alpha2  =   (sqrt(5d0)-1d0)/2d0
                d   =   bb-aa
                x1  =   aa+alpha1*d    
                x2  =   aa+alpha2*d    

                
                r_intertmp_ret = valfun_coup_ret(x1,share,deltas,deltac,gami,gamp,yi,yp,thetai,thetap,effi,effp, &
                                betac,betai,betap,pen,c0,&
                                vfc,vfi,vfp,vci,vcp,nu,psii,psip,SMF,wid_f,wid_m)
                f1  =   r_intertmp_ret(1)


                r_intertmp_ret = valfun_coup_ret(x2,share,deltas,deltac,gami,gamp,yi,yp,thetai,thetap,effi,effp, &
                                betac,betai,betap,pen,c0,&
                                vfc,vfi,vfp,vci,vcp,nu,psii,psip,SMF,wid_f,wid_m)
                f2  =   r_intertmp_ret(1)
                
                if(f1==-900000000000000000000000000d0 .AND. share == 0d0) write(*,*) 'vf is infinity -- couple', f2
                
                d   =   alpha1*alpha2*d

            do while(d>tol)
        
                d = d*alpha2
        
            if(f2<f1)then
        
                x2 = x1
                x1 = x1-d
                f2 = f1
                r_intertmp_ret = valfun_coup_ret(x1,share,deltas,deltac,gami,gamp,yi,yp,thetai,thetap,effi,effp, &
                                betac,betai,betap,pen,c0,&
                                vfc,vfi,vfp,vci,vcp,nu,psii,psip,SMF,wid_f,wid_m)
                f1  =   r_intertmp_ret(1)
            else
        
                x1 = x2
                x2 = x2+d
                f1 = f2
                r_intertmp_ret = valfun_coup_ret(x2,share,deltas,deltac,gami,gamp,yi,yp,thetai,thetap,effi,effp, &
                                betac,betai,betap,pen,c0,&
                                vfc,vfi,vfp,vci,vcp,nu,psii,psip,SMF,wid_f,wid_m)
                f2  =   r_intertmp_ret(1)
                
            end if
            
            end do    

            if(f2>f1) then
        
            x1=x2
            f1=f2
        
            end if
    
            
            ! store results in array
            
            r_intertmp_ret=valfun_coup_ret(x1,share,deltas,deltac,gami,gamp,yi,yp,thetai,thetap,effi,effp, &
                                betac,betai,betap,pen,c0,&
                                vfc,vfi,vfp,vci,vcp,nu,psii,psip,SMF,wid_f,wid_m)
            
            
            output_coup_ret(1) = x1
            output_coup_ret(2) = r_intertmp_ret(1)
            output_coup_ret(3) = r_intertmp_ret(2) 
            output_coup_ret(4) = r_intertmp_ret(3) 
            output_coup_ret(5) = r_intertmp_ret(4) 

        end function golden_coup_ret
        

end module
